#include<iostream>
#include<cstdlib>
#include<cstring>
#include "allocator_interface.h"
#include "memlib.h"
#include "assert.h"
#include <pthread.h>

/* All blocks must have a specified minimum alignment. */
#define ALIGNMENT 8

/* Rounds up to the nearest multiple of ALIGNMENT. */
#define ALIGN(size) (((size) + (ALIGNMENT-1)) & ~(ALIGNMENT-1))

/* The smallest aligned size that will hold a size_t value. */
#define SIZE_T_SIZE (ALIGN(sizeof(size_t)))

// Given that we're using a binned-free list, this defines the maximum number of bins.
// Sizes are exponential and in bytes, and on a 64-bit machine the max memory size
// request can be 2^64 bits, or 2^61 bytes. So, in theory we would require at most
// 61 bins. However, we get 256 bytes for global variables. So our array can have at
// most 256/(sizeof(pointer)) entries, or 31 elements. We could turn the array to
// a linked list to save on space too, but that seems unecessary.
// NUMBINS is a misnomer, since it describes the maximum index in our bins array.
// The actual number of bins is NUMBINS + 1
#define NUMBINS 30

namespace my
{
	// Global mutex lock
	pthread_mutex_t mtx;
	
	/* Misc. Auxiliary Functions */
	typedef uint8_t char_t;
	
	// Simple (and fairly quick) function that computes the floor of log2(x)
	inline size_t log2(size_t x){
		size_t r = 0;
		while(x >>= 1){	r++; }
		return r;
	}

	/* Memory block representation & associated functions */
	
	/* Represents a block of memory, and is the header of the memory block */
	struct MemBlock {		
		// Size of the memory block
		size_t head; 
		
		// Status of the block (so used or free)
		char_t status; 

		// Next and previous pointers for the linked list
		// These currently only work for when the block is binned
		// However, these could be extended to point to physically
		// contiguous blocks when the block is live. Can't think of
		// how that would help, though.
		MemBlock* prev;
		MemBlock* next;
	};
	
	void resizeMemBlock(MemBlock* block, size_t newBlockSize);
	void coalesceMemBlock(MemBlock* block);
	void removeMemBlock(MemBlock* block);
	
	// Simple enum for the memblock status.
	enum { USED, FREE };

	// Size of the header for a block of data.
	#define HEADER_SIZE (ALIGN(sizeof(MemBlock)))
	
	// Size of the smallest block of memory possible, which is just
	// the header size + the size of one MemBlock.
	#define MIN_BLOCK_SIZE (HEADER_SIZE + SIZE_T_SIZE)
	
	/* Pointer conversion macros */
	// Note: These macros aren't really type-safe, since we're toying around with pointers.
	
	// We cast our pointers to a char pointer so the compiler knows how many bytes to move when adding.
	// When converting pointers, we increase/decrease the pointer by HEADER_SIZE.
	
	// Given a pointer to a block of memory, gives us a pointer to the start of the memory space.
	#define MemBlock2mem(p) ((void*) ((char_t*)(p) + HEADER_SIZE))
	
	// Given a pointer to a memory spaces, gives us a pointer to the start of the block header.
	#define mem2MemBlock(mem) ((MemBlock*) ((char_t*)(mem) - HEADER_SIZE))
	
	// Given a pointer to a block of memory, gives us a pointer to the start of the next block.
	#define nextMemBlock(p) ((MemBlock*) ((char_t*)(p) + HEADER_SIZE + (p)->head))

	// These are the bins where we store our free MemBlocks
	static MemBlock* bins[NUMBINS];
	
	// Adds a block to the proper bin by taking the log of the size and doing a sorted insertion
	// The linked-list is therefore sorted by size, so we're doing a best-fit allocation when we
	// get blocks. Another option would be to sort them by memory location. On the one hand, it's
	// prone to creating more fragmentation. On the other hand, it might allow us to do better
	// coalescing since live blocks will be physically closer together.
	void binMemBlock(MemBlock* newMemBlock) {
		// First, we set the status to free.
		newMemBlock->status = FREE;
		// We set the prev and next pointers to null for now too.
		newMemBlock->prev=NULL;
		newMemBlock->next=NULL;

		// Next, we locate the right bin to add it to.
		size_t bin = log2(newMemBlock->head);
		if(bins[bin]==NULL){
			// If the bin is empty, we just straight up add it.
			bins[bin]=newMemBlock;
		} else if(newMemBlock->head < bins[bin]->head) {
			// If it's the smallest MemBlock, we place it at the start.
			newMemBlock->next = bins[bin];
			newMemBlock->next->prev = newMemBlock;
			bins[bin] = newMemBlock;
		} else {
			// Otherwise, we iterate through the linked list to find its location
			MemBlock * curr = bins[bin];
			while(curr->next != NULL && newMemBlock->head > curr->next->head) {
				curr=curr->next;
			}
			
			if(curr->next!=NULL) { 
				// If there's a block after us, we need to update its pointers.
				newMemBlock->next=curr->next;
				newMemBlock->next->prev=newMemBlock;
			}
			// We update our pointers and the ones of our predecessor.
			curr->next=newMemBlock;
			newMemBlock->prev=curr;
		}
	}

	/* Return a block for something that satisfies the block size request */
	MemBlock* getMemBlock(size_t size) {
		assert(size != 0);
		for (size_t i = log2(size); i < NUMBINS; i++){
			MemBlock* current = bins[i];
			while(current != NULL && current->head < size) {
				current = current->next;
			}
			
			if(current != NULL) {
				removeMemBlock(current);
				
				// Now that we have a block we can use, we grab it and we clear the next 
				// and previous pointers for safety and set the block's status to used.
				current->status = USED;
				current->next = NULL;
				current->prev = NULL;

				// We also resize the block to the correct size if we can				
				resizeMemBlock(current, size);
				
				return current;
			}
		}
		/* Returns NULL if no suitable sized block was found */
		return NULL;
	}
	
	// Simply removes a memblock from the binned free list by rearranging
	// the correct pointers appropriately
	void removeMemBlock(MemBlock* block){
		assert(block->status == FREE);
		assert(bins[log2(block->head)] != NULL);
		
		if(block->prev==NULL){
			assert(bins[log2(block->head)] == block);
			bins[log2(block->head)] = block->next;
		} else {
			block->prev->next = block->next;
		}
		
		if(block->next != NULL){
			block->next->prev = block->prev;
		}
	}

	// Attempts to reduce block down to newBlockSize if the difference is
	// large enough, and then breaks off and bins the extra space. If there isn't
	// enough extra space, the method does nothing.
	void resizeMemBlock(MemBlock* block, size_t newBlockSize){
		assert(newBlockSize == ALIGN(newBlockSize));
		assert(block->status == USED);
		assert(block->head >= newBlockSize);
		if (block->head >= newBlockSize + MIN_BLOCK_SIZE){
			MemBlock* newMemBlock = (MemBlock*)((char_t*)block + newBlockSize + HEADER_SIZE);
			newMemBlock->head = block->head - newBlockSize - HEADER_SIZE;
			block->head = newBlockSize;
			binMemBlock(newMemBlock);
		}
	}

	// Coalesces the entire memory.
	// This function works great once our bins start getting too small and we need to
	// bring them back up to a usable size.
	void coalesceFreeSpace(){
		for (MemBlock* curr = (MemBlock*) mem_heap_lo(); (char_t*) curr <= (char_t*) mem_heap_hi(); curr = nextMemBlock(curr)){
			if(curr->status == FREE){
				coalesceMemBlock(curr);
			}
		}
	}
	
	void coalesceSingleFreeBlock(MemBlock* block){
		assert(block->status == FREE);
		// While the next block is free and it falls within the bounds of the heap, we increase
		// the current block's size by the amout on the other block, remove it from our bins, and
		// repeat with the next block after that.
		MemBlock* nextBlock = nextMemBlock(block);
		if (((char_t*) nextBlock <= (char_t*) mem_heap_hi()) && (nextBlock->status == FREE)){
			// If the original block was free, we need to update its position in the bins.
			// Thus, we remove it from the bins now (just once)
			removeMemBlock(block);			
			removeMemBlock(nextBlock);
			block->head = block->head + nextBlock->head + HEADER_SIZE;
			binMemBlock(block);
		}
		
	}

	// Coalesces current block with other blocks if there is space. 
	// Meant to work with both active and freed blocks, so it only coalesces
	// forward. We could, in theory, coalesce backwards as well for free blocks.
	void coalesceMemBlock(MemBlock* block){
		bool removedFreeBlock = 0;
		
		// While the next block is free and it falls within the bounds of the heap, we increase
		// the current block's size by the amout on the other block, remove it from our bins, and
		// repeat with the next block after that.
		for (MemBlock* nextBlock = nextMemBlock(block); 
		((char_t*) nextBlock <= (char_t*) mem_heap_hi()) && (nextBlock->status == FREE);
		nextBlock = nextMemBlock(block)){
			// If the original block was free, we need to update its position in the bins.
			// Thus, we remove it from the bins now (just once)
			if (!removedFreeBlock && block->status == FREE){
				removeMemBlock(block);
				removedFreeBlock = 1;
			}
			
			removeMemBlock(nextBlock);
			block->head = block->head + nextBlock->head + HEADER_SIZE;
		}
		
		// If we removed the block at some point, we need to chuck it back into the bins.
		if (removedFreeBlock){
			binMemBlock(block);
		}
	}


	/*	Allocator Methods */

	/*
	* init - Initialize the malloc package.  Called once before any other
	* calls are made.  
	*/
	int allocator::init(){
		pthread_mutex_init(&mtx, NULL);
		for(int i = 0; i < NUMBINS; i++) {
			bins[i] = NULL;
		}
		return 0;
	}

	void * mm_malloc(size_t size);
	void mm_free(void *ptr);
	void * mm_realloc(void *ptr, size_t newsize);
	
	// malloc
	// First, we look through our bins to see if we can find a free block 
	// of the right size. If we find one, we return it. Otherwise, we have
	// to allocate more space on the heap.
	void * mm_malloc(size_t size){
		// We only really handle sizes that are aligned, so this saves some 
		// headaches in the long run.
		size = ALIGN(size);
		
		// We try to get a free block from our bins
		MemBlock* block = getMemBlock(size);
		if(block != NULL){
			// If we find one,  we grab it and return the corresponding
			// starting memory address.
			return MemBlock2mem(block);
		}
		
		/* We allocate a little bit of extra memory so that we can store the
		size and usage status */
		int blocksize = ALIGN(size + HEADER_SIZE);

		/* Expands the heap by the given number of bytes and returns a pointer to
		the newly-allocated area.  This is a slow call, so you will want to
		make sure you don't wind up calling it on every malloc. */
		void* p = mem_sbrk(blocksize);

		if ( p == (void*)-1 ){
			// Whoops, an error of some sort occurred.  We return NULL to let
			// the client code know that we weren't able to allocate memory.
			return NULL;
		} else {
			MemBlock* e = (MemBlock*) p;
			e->head = size;
			e->status = USED;
			e->prev = NULL;
			e->next = NULL;
			return MemBlock2mem(e);
		}
	}
	

	// free
	// We first try to coalesce the mem block as far as we can, and then
	// we place it into the correct bin.
	// Note that we don't defend against freeing pointers that have already
	// been freed before or invalid pointers.
	void mm_free(void *ptr){
		if (ptr == NULL) return;
		binMemBlock(mem2MemBlock(ptr));
	}	
	
	
	void * mm_realloc(void *ptr, size_t newsize){
		// We only really handle sizes that are aligned, so this saves some 
		// headaches in the long run.
		newsize = ALIGN(newsize);
		
		// Handling edge cases, when pointer is NULL or we're resizing to 0.
		if(ptr == NULL){
			return mm_malloc(newsize);
		} else if(newsize == 0){
			mm_free(ptr);
			return ptr;
		}
		
		MemBlock* block = mem2MemBlock(ptr);
		assert(block->status == USED);

		// Case 1: New size is smaller than the old size
		if (block->head >= newsize){
			// We attempt to reduce the block size, and then just return the original pointer
			resizeMemBlock(block, newsize);
			return ptr;
		}
		
		// Case 2: New size is bigger than the old size
		
		// First, we try to coalesce as far as we can. If we have more than enough to 
		// spare afterwards, we reduce a bit. Otherwise, we try to add more space if 
		// the current block is at the end of the heap.
		coalesceMemBlock(block);
		if (block->head >= newsize){
			resizeMemBlock(block, newsize);
			return ptr;
		} else if ((char_t*) nextMemBlock(block) > (char_t*) mem_heap_hi()){
			size_t difference = newsize - block->head;
			mem_sbrk(ALIGN(difference));
			block->head= block->head + ALIGN(difference);
			return ptr;
		}
		
		// If none of the above works, we take the old approach.
		void* newptr = mm_malloc(newsize);
		if(newptr==NULL){
			return NULL;
		}
		/* This is a standard library call that performs a simple memory copy. */
		memcpy( newptr, ptr, block->head );

		/* Release the old block. */
		mm_free(ptr);

		/* Return a pointer to the new block. */
		return newptr;
	}


	/*
	* check - This checks our invariant that the size_t header before every
	* block points to either the beginning of the next block, or the end of the
	* heap.
	*/
	int allocator::check(){
		// If nothing interesting has happened so far, we just return.
		if ((char_t*) mem_heap_lo() == (char_t*) mem_heap_hi()){
			return 0;
		}
		
		MemBlock* block = (MemBlock*) mem_heap_lo();
		while ((char_t*) block <= (char_t*) mem_heap_hi()){
			if (block->head != ALIGN(block->head)){
				// All blocks sizes must be aligned. The space isn't usable anyway, so
				// it's a reasonable restriction to enforce.
				return -1;
			}
			if (block->status == USED){
				// When a block is live, we check to see if next and prev pointers are null
				if (block->prev != NULL && block->next != NULL){
					// printf("ERROR: Used block maintains prev and next pointers \n");
					assert(false);
					return -1;
				}
				
				// If the block is being used, we iterate over all bins to ensure that it is not in there
				int bin = log2(block->head);
				for(MemBlock* current = bins[bin]; current != NULL; current = current-> next){
					if (current == block){
						// The used block is also in the free bins
						assert(false);
						return -1;
					}
				}

				
			} else if (block->status == FREE){
				// We ensure the block pointers are consistent.
				if (block->prev != NULL){
					if (block->prev->next != block){
						assert(false);
						return -1;
					}
				}
				
				if (block->next != NULL){
					if (block->next->prev != block){
						assert(false);
						return -1;
					}
				}
				
				// If the block is free, we iterate over all bins to ensure that it is
				// actually in the bin where it should.
				int bin = log2(block->head);
				for(MemBlock* current = bins[bin]; current != block; current = current-> next){
					if (current == NULL){
						// Our supposedly free block is not in the free bin list. Jolly.
						assert(false);
						return -1;
					}
				}
			} else {
				// Well wtf.
				assert(false);
				return -1;
			}
			
			block = nextMemBlock(block);
		}
		
		// If the last block doesn't point just outside the current heap
		// then something went horribly wrong
		if ((char_t*) block != (char_t*) mem_heap_hi() + 1){
			// printf("ERROR: Pointer outside last block doesn't match end of heap \n");
			assert(false);
			return -1;
		}		
		return 0;
	}

	void * allocator::malloc(size_t size){
		pthread_mutex_lock(&mtx);
		void * r = mm_malloc(size); 
		pthread_mutex_unlock(&mtx);
		return r;
	}
	
	void allocator::free(void *ptr){
		pthread_mutex_lock(&mtx);
		mm_free(ptr);
		pthread_mutex_unlock(&mtx);
	}
	
	// realloc
	void * allocator::realloc(void *ptr, size_t newsize){
		pthread_mutex_lock(&mtx);
		void * r = mm_realloc(ptr, newsize);
		pthread_mutex_unlock(&mtx);
		return r;
	}


	/*	Default Methods	*/

	/* call mem_reset_brk. */
	void allocator::reset_brk(){
		mem_reset_brk();
	}

	/* call mem_heap_lo */
	void * allocator::heap_lo(){
		return mem_heap_lo();
	}

	/* call mem_heap_hi */
	void * allocator::heap_hi(){
		return mem_heap_hi();
	}
};